package com.hulk.dryad.manage.framework.mdc.wrapper;


import com.alibaba.ttl.TtlCallable;
import com.alibaba.ttl.TtlRunnable;
import com.hulk.dryad.common.util.wrapper.CallableWrapper;
import com.hulk.dryad.common.util.wrapper.RunnableWrapper;
import com.hulk.dryad.manage.framework.mdc.constant.MdcConstants;
import com.hulk.dryad.manage.framework.mdc.util.TraceIdUtil;
import org.slf4j.MDC;

import java.util.concurrent.*;

/**
 * 线程池包装类
 *
 * @author hulk
 * @version 1.0
 * @Date 2020/03/26 15:29
 */
public class ThreadPoolExecutorMdcWrapper extends ThreadPoolExecutor {
    public ThreadPoolExecutorMdcWrapper(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit,
                                        BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    public ThreadPoolExecutorMdcWrapper(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit,
                                        BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
    }

    public ThreadPoolExecutorMdcWrapper(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit,
                                        BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, handler);
    }

    public ThreadPoolExecutorMdcWrapper(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit,
                                        BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory,
                                        RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    }

    @Override
    public void execute(Runnable task) {
		setMdcIdIfAbsent();
		//Runnable ttlRunnable = TtlRunnable.get(task);
        super.execute(new RunnableWrapper(task));
    }

    @Override
    public <T> Future<T> submit(Runnable task, T result) {
		setMdcIdIfAbsent();
		//Runnable ttlRunnable = TtlRunnable.get(task);
        return super.submit(new RunnableWrapper(task), result);
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
		setMdcIdIfAbsent();
		//Callable ttlCallable = TtlCallable.get(task);
        return super.submit(new CallableWrapper(task));
    }

    @Override
    public Future<?> submit(Runnable task) {
		setMdcIdIfAbsent();
		//Runnable ttlRunnable = TtlRunnable.get(task);
        return super.submit(new RunnableWrapper(task));
    }

	private   void setMdcIdIfAbsent() {
		if (MDC.get(MdcConstants.TRACE_ID) == null) {
			MDC.put(MdcConstants.TRACE_ID, TraceIdUtil.getTraceId());
		}
	}
}
